import logging
import fvcore.nn.weight_init as weight_init
from typing import Optional
import torch
from torch import nn, Tensor
from torch.nn import functional as F
import numpy as np

from detectron2.config import configurable
from detectron2.layers import Conv2d
from detectron2.structures import Boxes, ImageList, Instances, BitMasks


from .position_encoding import PositionEmbeddingSine
from .maskformer_transformer_decoder import TRANSFORMER_DECODER_REGISTRY
from .mix_decoder import MixDecoder

class RetrievalLocalAttention(nn.Module):
    # @configurable
    def __init__(
        self,
        out_channels = 1024,
        align_size = 16
    ):
        super().__init__()
        in_channels = 256
        self.mix_decoder = MixDecoder(
            in_channels=in_channels,
            in_h = align_size,
            in_w = align_size,
            out_channels= out_channels // 4,
            mix_depth=4,
            mlp_ratio=1,
            out_rows=4,
        )
        self.out_feature_dim = out_channels
        self.align_size = align_size
        
    ##in multi_scale_features B C H W 
    '''
    EXAMPLE
    SRC IMAGE 10*512*512
    features:res2 10, 256, 128, 128
    features:res2 10, 512, 64, 64
    features:res2 10, 1024, 32, 32
    features:res2 10, 2048, 16, 16
    mask_features :10, 256, 128, 128
    transformer_encoder_features :10, 256, 16, 16
    multi_scale_features[0]:10, 256, 16, 16
    multi_scale_features[1]:10, 256, 32, 32
    multi_scale_features[2]:10, 256, 64, 64
    
    b_id_boxes:
        int tensor:b,id,x0,y0,x1,y1 
    '''
    def inference(self,feature,bboxes,qfeatures=None):
        
        boxex_features = torch.tensor([]).to(feature)
        h16 = self.align_size
        w16 = self.align_size
        
        for box in bboxes:
            mask_feature_box = feature[...,box[0]:box[2],box[1]:box[3]]
            mask_feature_box = F.interpolate(
                mask_feature_box, size=(h16,w16), mode="bilinear", align_corners=False
            )
            boxex_features = torch.cat((boxex_features,mask_feature_box),dim=0)     
                   
        mix_feats = self.mix_decoder.forward(boxex_features)
        return mix_feats